import matplotlib.pyplot as plt
import numpy as np
from matplotlib import cm
from matplotlib.animation import FuncAnimation


class MultiYAxisPlot:
    def __init__(self, x, y_data, labels, colors=None, y_lims=None):
        """
        初始化多Y轴折线图。

        参数:
            x: x轴数据（初始值）
            y_data: 一个列表，包含每个Y轴的初始数据
            labels: 一个列表，包含每个Y轴的标签
            colors: 一个列表，包含每个Y轴的线条颜色（可选）
            y_lims: 一个列表，包含每个Y轴的范围（可选）
        """
        if len(y_data) < 2:
            raise ValueError("至少需要两个Y轴的数据！")

        # 设置中文字体和负号显示
        plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置中文字体为黑体
        plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

        self.x = x
        self.y_data = y_data
        self.labels = labels
        self.colors = self.generate_color_array(len(y_data))
        # self.colors = colors if colors else plt.cm.tab10(np.linspace(0, 1, len(y_data)))
        self.y_lims = y_lims if y_lims else [None] * len(y_data)

        # 创建主图和第一个Y轴
        self.fig, self.axs = plt.subplots(figsize=(10, 6))  # 调整图形大小

        self.lines = []

        # 创建第一个Y轴的折线
        self.lines.append(self.axs.plot(x, y_data[0], linestyle='--', color=self.colors[0],
                                        alpha=0.8, lw=3, label=labels[0])[0])
        self.axs.set_xlabel("里程 m")
        self.axs.set_ylabel(labels[0], color=self.colors[0])
        if self.y_lims[0]:
            self.axs.set_ylim(self.y_lims[0])

        # 创建额外的Y轴
        self.extra_axes = []
        for i in range(1, len(y_data)):
            ax = self.axs.twinx()
            if i > 1:
                # 将额外的Y轴移动到右侧合适的位置
                ax.spines['right'].set_position(('outward', 50 * (i - 1)))
            line = ax.plot(x, y_data[i], color=self.colors[i], lw=2, label=labels[i])[0]
            ax.set_ylabel(labels[i], color=self.colors[i])
            if self.y_lims[i]:
                ax.set_ylim(self.y_lims[i])
            self.lines.append(line)
            self.extra_axes.append(ax)

        # 设置图例
        self.fig.legend(
            loc="lower center",  # 将图例放在底部中央
            bbox_to_anchor=(0.5, 0),  # 图例的位置 (x=0.5，图形正下方)
            fontsize="small",  # 设置字体大小
            ncol=len(self.labels)  # 图例横向排列，列数等于图例数量
        )

        # 调整布局
        plt.subplots_adjust(bottom=0.15)  # 留出足够的底部空间

        # 自动调整布局
        plt.tight_layout()

    def setTitle(self, title):
        # 动态更新标题
        plt.title(title, fontsize=16)

    def update(self, new_y_data):
        """
        更新折线图的数据。

        参数:
            new_y_data: 一个列表，包含每个Y轴的新数据
        """
        if len(new_y_data) != len(self.lines):
            raise ValueError("新数据的长度必须与初始化时一致！")

        for i, line in enumerate(self.lines):
            # 获取当前数据
            new_data = new_y_data[i]

            # 动态调整 Y 轴范围
            if i == 0:
                self.axs.set_ylim(np.min(new_data), np.max(new_data))
            else:
                self.extra_axes[i-1].set_ylim(np.min(new_data), np.max(new_data))

            # 更新数据
            line.set_ydata(new_data)



        # 动态调整Y轴范围（可选）
        for i, ax in enumerate([self.axs] + self.extra_axes):
            ax.relim()
            ax.autoscale_view()

        # 更新绘图
        # plt.draw()
        # 更新绘图
        self.fig.canvas.draw_idle()
        self.fig.canvas.flush_events()
        plt.pause(0.00001)

    def animate(self, update_func, frames=100, interval=200):
        """
        创建动画，并通过update_func动态更新数据。

        参数:
            update_func: 一个函数，用于生成新的数据。每次调用应返回一个包含所有Y轴新数据的列表。
            frames: 动画的总帧数
            interval: 每帧的间隔时间（毫秒）
        """
        ani = FuncAnimation(self.fig, lambda frame: self.update(update_func(frame)),
                            frames=range(frames), interval=interval, blit=False)
        plt.show()

    def generate_color_array(self, num, colormap_name='viridis'):
        """
        生成一个包含指定数量颜色的数组。

        参数:
            num_colors: 需要生成的颜色数量
            colormap_name: 使用的颜色映射名称（如 'viridis', 'plasma', 'cool', 等）
        返回:
            colors: 一个形状为 (num_colors, 4) 的数组，包含 RGBA 值
        """
        # 获取颜色映射对象
        colormap = cm.get_cmap(colormap_name, num)

        # 生成颜色数组
        colors2 = colormap(np.linspace(0, 1, num))

        return colors2


# 示例用法
if __name__ == "__main__":
    # 初始数据
    x = np.linspace(0, 10, 100)
    y1 = np.sin(x)
    y2 = np.cos(x) * 2
    y3 = np.tan(x) / 10
    y4 = np.sin(x) * np.cos(x)
    y5 = np.sin(x) + np.cos(x)

    labels = ["Y1", "Y2", "Y3", "Y4", "Y5"]
    colors = ["blue", "orange", "green", "red", "purple"]

    # 创建多Y轴折线图
    plot = MultiYAxisPlot(x, [y1, y2, y3, y4, y5], labels)


    # 定义一个函数，用于动态生成新的数据
    def update_func(frame):
        # 模拟动态数据更新
        new_y1 = np.sin(x + frame * 0.1)
        new_y2 = np.cos(x + frame * 0.1) * 2
        new_y3 = np.tan(x + frame * 0.1) / 10
        new_y4 = np.sin(x + frame * 0.1) * np.cos(x + frame * 0.1)
        new_y5 = np.sin(x + frame * 0.1) + np.cos(x + frame * 0.1)
        return [new_y1, new_y2, new_y3, new_y4, new_y5]


    # 启动动画
    plot.animate(update_func, frames=100, interval=200)
